#!/bin/bash
# HunyuanOCR NPU 部署脚本
# 模型路径: /data0/HunyuanOCR

set -e

# 配置
MODEL_PATH="/data0/HunyuanOCR"
CONTAINER_NAME="hunyuan-ocr-api"
IMAGE_NAME="hunyuan-ocr-npu:latest"
HOST_PORT=18001
CONTAINER_PORT=8000
NPU_DEVICE="davinci0"  # 可根据需要修改: davinci0, davinci1, davinci2, davinci3

echo "=========================================="
echo "HunyuanOCR NPU 部署脚本"
echo "=========================================="

# 检查Ascend CANN环境
echo "检查Ascend CANN环境..."
ASCEND_DRIVER="/usr/local/Ascend/driver"
ASCEND_DRIVER_LIB="/usr/local/Ascend/driver/lib64"
ASCEND_DRIVER_VERSION="/usr/local/Ascend/driver/version.info"
ASCEND_TOOLKIT="/usr/local/Ascend/ascend-toolkit"
ASCEND_INSTALL_INFO="/etc/ascend_install.info"
DCMI_PATH="/usr/local/dcmi"
NPU_SMI="/usr/local/bin/npu-smi"

# 检查NPU设备
echo "检查NPU设备..."
npu-smi info || echo "警告: 无法执行npu-smi，请确保已安装昇腾驱动"

# 检查关键路径和文件
MISSING_PATHS=()

if [ ! -d "$ASCEND_DRIVER" ]; then
    echo "警告: 未找到 $ASCEND_DRIVER 目录"
    MISSING_PATHS+=("$ASCEND_DRIVER")
else
    echo "✓ 找到驱动目录: $ASCEND_DRIVER"
    if [ -d "$ASCEND_DRIVER_LIB" ]; then
        if find "$ASCEND_DRIVER_LIB" -name "libascend_hal.so" 2>/dev/null | head -1 | grep -q .; then
            echo "✓ 找到 libascend_hal.so"
        else
            echo "警告: 在 $ASCEND_DRIVER_LIB 中未找到 libascend_hal.so"
        fi
    fi
fi

if [ ! -d "$ASCEND_TOOLKIT" ]; then
    echo "提示: 未找到 $ASCEND_TOOLKIT 目录（可选，但推荐挂载）"
else
    echo "✓ 找到CANN工具包目录: $ASCEND_TOOLKIT"
fi

if [ ! -f "$ASCEND_DRIVER_VERSION" ]; then
    echo "警告: 未找到 $ASCEND_DRIVER_VERSION 文件"
    MISSING_PATHS+=("$ASCEND_DRIVER_VERSION")
else
    echo "✓ 找到驱动版本信息: $ASCEND_DRIVER_VERSION"
fi

if [ ! -f "$ASCEND_INSTALL_INFO" ]; then
    echo "警告: 未找到 $ASCEND_INSTALL_INFO 文件"
    MISSING_PATHS+=("$ASCEND_INSTALL_INFO")
else
    echo "✓ 找到安装信息: $ASCEND_INSTALL_INFO"
fi

if [ ! -d "$DCMI_PATH" ]; then
    echo "提示: 未找到 $DCMI_PATH 目录（可选）"
else
    echo "✓ 找到DCMI路径: $DCMI_PATH"
fi

if [ ! -f "$NPU_SMI" ]; then
    echo "提示: 未找到 $NPU_SMI 工具（可选）"
else
    echo "✓ 找到 npu-smi 工具: $NPU_SMI"
fi

if [ ${#MISSING_PATHS[@]} -gt 0 ]; then
    echo ""
    echo "警告: 以下路径未找到，容器可能无法正常访问NPU运行时库："
    for path in "${MISSING_PATHS[@]}"; do
        echo "  - $path"
    done
    echo "这可能导致 libascend_hal.so 错误"
    echo ""
fi

# 检查模型目录
if [ ! -d "$MODEL_PATH" ]; then
    echo "错误: 模型目录 $MODEL_PATH 不存在"
    echo "请先下载模型:"
    echo "  huggingface-cli download tencent/HunyuanOCR --local-dir $MODEL_PATH"
    exit 1
else
    echo "✓ 找到模型目录: $MODEL_PATH"
fi

# 构建Docker镜像
echo ""
echo "构建Docker镜像..."
cd "$(dirname "$0")/HunyuanOCR-Ascend"
docker build -t ${IMAGE_NAME} .
cd ..

# 停止并删除旧容器
echo ""
echo "停止旧容器..."
docker stop ${CONTAINER_NAME} 2>/dev/null || true
docker rm ${CONTAINER_NAME} 2>/dev/null || true

# 启动新容器 (参考 MiniCPM 部署方式)
echo ""
echo "启动 HunyuanOCR 服务 (NPU: ${NPU_DEVICE}, 端口: ${HOST_PORT})..."

# 挂载所有 NPU 设备（驱动初始化需要）
NPU_DEVICE_ARGS=""
for dev in /dev/davinci*; do
    if [ -e "$dev" ]; then
        NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=$dev"
    fi
done
if [ -e "/dev/davinci_manager" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/davinci_manager"
fi
if [ -e "/dev/devmm_svm" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/devmm_svm"
fi
if [ -e "/dev/hisi_hdc" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/hisi_hdc"
fi

docker run -d \
  --name ${CONTAINER_NAME} \
  --restart=unless-stopped \
  --privileged \
  -p ${HOST_PORT}:${CONTAINER_PORT} \
  $NPU_DEVICE_ARGS \
  -e ASCEND_VISIBLE_DEVICES=0 \
  -e ASCEND_RT_VISIBLE_DEVICES=0 \
  -e ASCEND_DEVICE_ID=0 \
  -v /usr/local/dcmi:/usr/local/dcmi \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
  -v /etc/ascend_install.info:/etc/ascend_install.info \
  -v /data0:/data0:ro \
  -e MODEL_PATH=${MODEL_PATH} \
  -e USE_BF16=true \
  ${IMAGE_NAME}

# 等待服务启动
echo ""
echo "等待服务启动..."
sleep 5

# 检查服务状态
echo ""
echo "=========================================="
echo "部署完成！"
echo "=========================================="
echo ""
echo "服务信息:"
echo "  容器名称: ${CONTAINER_NAME}"
echo "  NPU设备:  ${NPU_DEVICE}"
echo "  端口:     ${HOST_PORT}"
echo "  模型路径: ${MODEL_PATH}"
echo ""
echo "服务状态:"
docker ps | grep ${CONTAINER_NAME} || echo "警告: 容器可能未成功启动"
echo ""
echo "API地址:"
echo "  健康检查:    http://localhost:${HOST_PORT}/health"
echo "  API文档:     http://localhost:${HOST_PORT}/docs"
echo "  OCR接口:     http://localhost:${HOST_PORT}/v1/ocr"
echo "  聊天接口:    http://localhost:${HOST_PORT}/v1/chat/completions"
echo ""
echo "测试命令:"
echo "  curl http://localhost:${HOST_PORT}/health"
echo ""
echo "查看日志:"
echo "  docker logs -f ${CONTAINER_NAME}"
echo ""
echo "查看NPU使用情况:"
echo "  npu-smi info"

